# -*- coding: utf-8 -*-
"""
@Time ： 2020-12-08 17:19
@Auth ： liangpw3
@Description：

"""
from sklearn import linear_model

from algo.Algo_interface import Algo_interface
from sklearn.neighbors import KNeighborsClassifier
import lightgbm as lgb
from sklearn.ensemble import RandomForestClassifier
from sklearn import svm
from xgboost.sklearn import XGBClassifier
from sklearn.ensemble import GradientBoostingClassifier
from sklearn import linear_model

class AudioClassifier(Algo_interface):
    def __init__(self, model_type, model_name, model_params):
        self.task_type = model_type
        self.model_name = model_name
        self.model_params = model_params
        self.model = None
        self.build_model()
        # return self.model

    def set_model(self, model):
        self.model = model
        return 1

    def get_model(self):
        return self.model

    def build_model(self):
        if self.model_name == '逻辑回归':
            self.model = linear_model.LogisticRegression(**self.model_params)
        elif self.model_name == 'svm':
            self.model = svm.SVC(**self.model_params)
        elif self.model_name == '随机森林':
            self.model = RandomForestClassifier(**self.model_params)
        elif self.model_name == 'lgb':
            self.model = lgb.LGBMClassifier(**self.model_params)
        elif self.model_name == 'xgboost':
            self.model = XGBClassifier(**self.model_params)
        elif self.model_name == 'gbdt':
            self.model = GradientBoostingClassifier(**self.model_params)
        elif self.model_name == 'knn':
            self.model = KNeighborsClassifier(**self.model_params)
    def train(self,data):
        x_train, y_train=data
        self.model.fit(x_train, y_train)
        return 1
    def predict(self,data):
        y_pred = self.model.predict(data)
        return y_pred